from sentence_transformers import SentenceTransformer
import chromadb
from llm.doubao import doubao_qa
from llm.local import ollama_qa
import json

model = SentenceTransformer("all-MiniLM-L6-v2")

client = chromadb.PersistentClient(path="./chromadb_data")
collection = client.get_or_create_collection("rag")


def get_query_embedding(query):
    return model.encode(query).tolist()


def retrieve_related_chunks(query_embedding, n_results=3):
    results = collection.query(query_embeddings=[query_embedding], n_results=n_results)
    related_chunks = results.get("documents")
    if not related_chunks or not related_chunks[0]:
        exit(1)
    return related_chunks[0]


def rewrite_query_multiple_versions(query):
    rewriting_prompt = f"""
    任务：请对以下的用户查询进行多角度重写，生成3到5个不同表达方式的查询版本

    要求：
    1. 保持原始的查询核心语义不变
    2. 使用不同的表达方式和措辞
    3. 可以补充相关的关键字或概念
    4. 考虑不同的查询意图和表达习惯
    5. 确保每个重写版本都是完整清晰的查询


    原始查询 ：{query}

    以JSON格式返回：
    {{
    "rewritten_queries": [
        "重写版本1",
        "重写版本2",
        "重写版本3"
    ]
    }}
    """
    decomposition_result = ollama_qa(rewriting_prompt)
    print("decomposition_result", decomposition_result)
    jsons_tart = decomposition_result.find("{")
    jsons_end = decomposition_result.rfind("}") + 1
    json_str = decomposition_result[jsons_tart:jsons_end]
    parsed_result = json.loads(json_str)
    rewritten_queries = parsed_result.get("rewritten_queries", [])
    print("rewritten_queries", rewritten_queries)
    return rewritten_queries


def parallel_retrieve_multiple_queries(rewritten_queries, n_results_per_query=2):
    all_related_chunks = []
    for i, query in enumerate(rewritten_queries):
        query_embedding = get_query_embedding(query)
        related_chunks = retrieve_related_chunks(query_embedding, n_results_per_query)
        all_related_chunks.extend(related_chunks)
    unique_chunks = list(set(all_related_chunks))
    return unique_chunks


def integrate_answers(original_query, all_related_chunks):
    context = "\n".join(all_related_chunks)
    integration_prompt = f"""
基于以下检索到的相关信息，请回答用户的原始问题

检索到的相关信息为:
{context}

用户的原始问题为:
{original_query}

请提供一个全面准确的回答，确保
1. 回答覆盖原问题的所有的方面，不要遗漏
2. 基于检索到的信息进行回答
3. 如果信息不足，则明确指出
4. 保持逻辑清晰，结构合理
"""
    final_answer = ollama_qa(integration_prompt)
    return final_answer


if __name__ == "__main__":
    query = input("请输入你的问题:")
    # 多角度查询重写
    rewritten_queries = rewrite_query_multiple_versions(query)
    print("rewritten_queries:", rewritten_queries)
    # 并行检索多个改写后的查询
    all_related_chunks = parallel_retrieve_multiple_queries(
        rewritten_queries, n_results_per_query=2
    )
    # 整合答案
    final_answer = integrate_answers(query, all_related_chunks)
    print(f"answer:{final_answer}")
